# -*- coding: utf-8 -*-
from selenium import webdriver
from lxml import etree
from fake_useragent import UserAgent
import os
import csv
import time
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC

headers = {
    "User-Agent": UserAgent().random
}


def write_to_csv(save_path, filename, items):  # 保存为CSV文件
    if not os.path.exists(save_path):  # 判断文件路径是否存在，若不在，则创建
        os.mkdir(save_path)
    with open(save_path + '/' + filename + '.csv', 'a+', newline='') as f:
        fieldnames = ['股票代码', '股票名称', '数据地址', '今日价格', '今日排名', '今日涨跌', '所属模块']
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        try:
            writer.writerows(items)
        except:
            pass


def Page_Level(save_path, filename, myPage):  # 一级页面：频道名称和频道url
    dom = etree.HTML(myPage)
    items = []
    stock_code = dom.xpath('//*[@id="dt_1"]/tbody/tr/td[2]/a/text()')
    stock_name = dom.xpath('//*[@id="dt_1"]/tbody/tr/td[3]/a/text()')
    stock_data_src = dom.xpath('//*[@id="dt_1"]/tbody/tr/td[4]/a[2]/@href')
    stock_new_price = dom.xpath('//*[@id="dt_1"]/tbody/tr/td[5]/span/text()')
    stock_new_rank = dom.xpath('//*[@id="dt_1"]/tbody/tr/td[7]/span/text()')
    stock_new_float = dom.xpath('//*[@id="dt_1"]/tbody/tr/td[8]/span/text()')
    stock_module = dom.xpath('//*[@id="dt_1"]/tbody/tr/td[15]/a/text()')

    for s_code, s_name, sd_src, sn_price, sn_rank, sn_float, s_module in zip(stock_code, stock_name, stock_data_src,
                                                                             stock_new_price, stock_new_rank,
                                                                             stock_new_float, stock_module):
        item = {}
        item['股票代码'] = s_code
        item['股票名称'] = s_name
        item['数据地址'] = sd_src
        item['今日价格'] = sn_price
        item['今日排名'] = sn_rank
        item['今日涨跌'] = sn_float
        item['所属模块'] = s_module
        items.append(item)
    write_to_csv(save_path, filename, items)


def spider(myPage, title):
    save_path = '下载的文件'  # 保存内容
    filename = title
    Page_Level(save_path, filename, myPage)  # 爬取


def main():
    # 加载驱动
    driver = webdriver.Firefox()
    driver.get('http://data.eastmoney.com/zjlx/list.html')
    name = "股票"
    # 选择股票类型
    print("选择股票类型：\n")
    choice = int(input("1.全部股票  2.沪深A股  3.沪市A股  4.科创板  5.深市A股  6.创业板  7.中小板  8.沪市B股  9.深市B股  \n"))
    element = WebDriverWait(driver, 5, 0.5).until(EC.presence_of_element_located((By.ID, "mk_type")))  # 显式等待
    if choice == 1:
        element.find_element_by_xpath('//*[@id="mk_type"]/li[1]').click()
        name = "全部股票"
    elif choice == 2:
        element.find_element_by_xpath('//*[@id="mk_type"]/li[2]').click()
        name = "沪深A股"
    elif choice == 3:
        element.find_element_by_xpath('//*[@id="mk_type"]/li[3]').click()
        name = "沪市A股"
    elif choice == 4:
        element.find_element_by_xpath('//*[@id="mk_type"]/li[4]').click()  #
        name = "科创板"
    elif choice == 5:
        element.find_element_by_xpath('//*[@id="mk_type"]/li[5]').click()
        name = "深市A股"
    elif choice == 6:
        element.find_element_by_xpath('//*[@id="mk_type"]/li[6]').click()
        name = "创业板"
    elif choice == 7:
        element.find_element_by_xpath('//*[@id="mk_type"]/li[7]').click()
        name = "中小板"
    elif choice == 8:
        element.find_element_by_xpath('//*[@id="mk_type"]/li[8]').click()  #
        name = "沪市B股"
    elif choice == 9:
        element.find_element_by_xpath('//*[@id="mk_type"]/li[9]').click()  #
        name = "深市B股"
    else:
        print("\n输入有误\n")
    # 输入爬取的页数
    if not (choice == 4 or choice == 8 or choice == 9):
        page = int(input("\n请输入要获取的数据的页数: "))
    else:
        page = 1
    for i in range(1, page + 1):
        html = driver.page_source
        try:
            spider(html, name + str(i))
            driver.find_element_by_xpath('//*[@id="PageCont"]/a[8]').click()
            time.sleep(3)
        except Exception:
            print("只有这一页")
    driver.quit()


if __name__ == "__main__":
    main()
